# _*_ coding: utf-8 _*_
"""
@ 时间    ：2024/10/22 19:10
@ 作者    ：旺财
@ 文件    ：02 一元二次模型.py
@ 说明    ：预测工龄与薪水的关系
"""
import pandas
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
from matplotlib import pyplot as plt
from sklearn.metrics import r2_score

# 1.读取数据
df = pandas.read_excel('IT行业收入表.xlsx')
X = df[['工龄']].values.reshape(-1, 1)  # 自变量需要二维数据结构
Y = df['薪水']

# 2.模型训练
poly_reg = PolynomialFeatures(degree=2)
X_ = poly_reg.fit_transform(X)
regr = LinearRegression()
regr.fit(X_, Y)

# 3.方程构造
_, b, a = [int(x) for x in regr.coef_]  # b a分别为一次项与二次项系数
c = int(regr.intercept_)  # 常数项
equation = f'方程: y = {a}x^2 + {b}x + {c}'.replace('+ -', '- ')


# 4.模型预测
x = 5.4
y = regr.predict(poly_reg.transform([[x]]))[0]
print(f'模型预测: IT行业中,工龄{x}年的人薪水预测结果为:{int(y)}元')

# 5.模型评估
y_pred = regr.predict(X_)
r2 = r2_score(Y, y_pred)
print(f'模型评估: 模型的预期准确度为:{round(r2 * 100, 2)}%')

# 6.模型可视化
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.scatter(X, Y)
plt.plot(X, regr.predict(X_), color='red', label=equation)
plt.xlabel('工龄')
plt.ylabel('薪水')
plt.legend(loc='upper left')
plt.show()
